from ..base import *

class StatsPublic():
	'''
	统计基础类
	'''
	@staticmethod
	def calc_on_array_or_single(calc_func: Callable, x: FloatOrArray) -> FloatOrArray:
		'''
		对数组或单个值进行计算
		'''
		if isinstance(x, np.ndarray) or isinstance(x, list):
			return np.array([calc_func(x) for x in x])
		return calc_func(x)
	
	@staticmethod
	def scale_and_offset(x:FloatOrArray, option: tuple[float, float])->FloatOrArray:
		return StatsPublic.calc_on_array_or_single(lambda x: x * option[1] + option[0] if option is not None else x, x)
	
	@staticmethod
	def accum_expectation(e1: float, e2: float, n1: float, n2: float) -> float:
		'''
		期望的合并计算公式
		'''
		if np.abs(n1 + n2) < MinFloat:
			return 0
		
		return (e1 * n1 + e2 * n2)/(n1 + n2)

	@staticmethod
	def calc_median(mean: float, stddev: float, skew: float) -> float:
		"""计算中位数"""
		return mean - skew * stddev / 3

	@staticmethod
	def calc_mode(mean: float, stddev: float, skew: float) -> float:
		"""计算众数"""
		return mean - 3*(mean - StatsPublic.calc_median(mean, stddev, skew))

	@staticmethod
	def pdf_from_cf(cf, x_min=-10, x_max=10, n_points=1000):
		'''
		通过特征函数生成PDF
		'''
		# 离散化 t 轴
		t = np.linspace(-50, 50, n_points)
		dt = t[1] - t[0]  # 步长
		
		# 计算特征函数在 t 上的值
		phi_t = cf(t)
		
		# 计算逆傅里叶变换
		x = np.linspace(x_min, x_max, n_points)
		f_x = np.fft.fftshift(np.fft.ifft(np.fft.ifftshift(phi_t))) * n_points * dt / (2 * np.pi)
		
		return x, np.abs(np.real(f_x))  # 取实部（数值误差可能导致虚部很小）

class Distribution():
	'''抽象类'''
	def __init__(self, lower_bound: float = None, upper_bound: float = None):  # type: ignore
		super().__init__()
		'''
		初始化分布类
		lower_bound: float, 分布的下限
		upper_bound: float, 分布的上限
		group_count: int, 分组数量
		'''
		self.lower_bound = lower_bound # 下边界
		self.upper_bound = upper_bound # 上边界
		self.n = 0 # 样本数
		self.expect0 = 0 # 零阶矩
		self.expect1 = 0 # 一阶矩
		self.expect2 = 0 # 二阶矩
		self.expect3 = 0 # 三阶矩
		self.expect4 = 0 # 四阶矩
		self.mean = 0 if self.lower_bound == None or self.upper_bound == None else (self.lower_bound + self.upper_bound) / 2 # 均值
		self.mode = 0 if self.lower_bound == None or self.upper_bound == None else (self.lower_bound + self.upper_bound) / 2 # 众数
		self.median = 0 if self.lower_bound == None or self.upper_bound == None else (self.lower_bound + self.upper_bound) / 2 # 中位数
		self.variance = 0 # 方差
		self.stddev = 0 # 标准差
		self.skewness = 0 # 偏度
		self.hurst = 0 # 峰度
		self.min = None # 最小值
		self.max = None # 最大值
		self.N = 0 # 总样本数
		self.prob = 0 # 概率
		self.entropy = 0 # 熵

	@property
	def x_range(self)->tuple[float|None, float|None]: # type: ignore
		'''
		获取分布的下限和上限
		'''
		return self.min if self.lower_bound == None else self.lower_bound, self.max if self.upper_bound == None else self.upper_bound
	
	def append_data(self, x: float, n: float = 1, N: float = None): # type: ignore
		'''
		添加数据
		params:
			x: float, 数据点
			n: float, 数据个数
		'''
		if self.lower_bound is not None and x < self.lower_bound:
			return
		if self.upper_bound is not None and x > self.upper_bound:
			return
		# 更新最小值和最大值
		self.min = x if self.min == None else min(self.min, x)
		self.max = x if self.max == None else max(self.max, x)
		# 更新期望
		N = self.n + n if N == None else N
		self.N = N
		self.expect0 = self.n + n
		self.expect1 = StatsPublic.accum_expectation(self.expect1, x, self.n, n)
		self.expect2 = StatsPublic.accum_expectation(self.expect2, x ** 2, self.n, n)
		self.expect3 = StatsPublic.accum_expectation(self.expect3, x ** 3, self.n, n)
		self.expect4 = StatsPublic.accum_expectation(self.expect4, x ** 4, self.n, n)
		self.n += n
		self.prob = self.n / N
		self.entropy = self.prob * np.log2(self.prob)
		self._update_parameters()

	def get_info(self, label: str)->str:
		'''
		打印信息
		params:
			label: str, 标签
		'''
		min, max = self.x_range
		if self.n == 0:
			return f"[{label},n={self.n:.3f},min={min:.3f},max={max:.3f}]"
		else:

			return f"[{label},n={self.n:.3f},μ={self.mean:.3f},σ={self.stddev:.3f},γ={self.skewness:.3f},δ={self.hurst:.3f},δ={self.hurst:.3f},min={min:.3f},max={max:.3f}]"

	T = TypeVar('T', bound='Distribution')
	def combine(self: T, dist: T|Iterable[T])-> T:
		'''
		合并两/多个分布
		params:
			dist: Distribution, 待合并的分布
		'''
		print('----开始合并----')
		if isinstance(dist, Iterable):
			for d in dist:
				self.combine(d)
			return self
		
		# print(self.info('原分布'))
		# print(dist.info('加分布'))
		self.min = min(self.min if self.min != None else np.inf, dist.min if dist.min != None else np.inf)
		self.max = max(self.max if self.max != None else -np.inf, dist.max if dist.max != None else -np.inf)
		self.lower_bound = min(self.lower_bound if self.lower_bound != None else np.inf, dist.lower_bound if dist.lower_bound != None else np.inf)
		self.upper_bound = max(self.upper_bound if self.upper_bound != None else -np.inf, dist.upper_bound if dist.upper_bound != None else -np.inf)
		# 更新期望
		self.expect0 = StatsPublic.accum_expectation(self.expect0, dist.expect0, self.n, dist.n)
		self.expect1 = StatsPublic.accum_expectation(self.expect1, dist.expect1, self.n, dist.n)
		self.expect2 = StatsPublic.accum_expectation(self.expect2, dist.expect2, self.n, dist.n)
		self.expect3 = StatsPublic.accum_expectation(self.expect3, dist.expect3, self.n, dist.n)
		self.expect4 = StatsPublic.accum_expectation(self.expect4, dist.expect4, self.n, dist.n)
		self.n += dist.n
		self._update_parameters()

		# print(self.info('新分布'))
		print('----完成合并----')
		return self # type: ignore

	@property
	def x_range(self)->tuple[float|None, float|None]:
		'''
		获取数据范围, 如果有期望范围，则取期望范围，否则取最小值和最大值
		returns: tuple[float, float], 数据范围
		'''
		return self.min if self.lower_bound == None else self.lower_bound, self.max if self.upper_bound == None else self.upper_bound

	def calc_deviation(self, x: FloatOrArray, type: str = 'mean')->FloatOrArray:
		'''
		计算离散度
		params:
			x: float, 待计算的点
			type: str, 计算类型，包括均值、中位数、众数
		returns: float, 离散度
		'''
		expect_value = self.mean
		if type == 'median':
			expect_value = self.median
		elif type =='mode':
			expect_value = self.mode
		return StatsPublic.calc_on_array_or_single(lambda x: np.abs(x - expect_value) / self.stddev if self.stddev > MinFloat else 0, x)	
	
	def in_range(self, x: float)->bool:
		'''
		判断数据是否在范围内
		params:
			x: float, 待判断的数据
		returns: bool, 是否在范围内
		'''
		return self.lower_bound <= x and x <= self.upper_bound

	def _update_parameters(self):
		'''
		更新参数
		'''
		self.mean = self.expect1
		# 更新方差
		self.variance = self.expect2 - self.expect1 ** 2
		# 更新标准差
		self.variance = max(self.variance, 0)
		self.stddev = np.sqrt(self.variance)

		# 更新偏度
		# 更新峰度
		if self.stddev >= MinFloat:
			self.skewness = (self.expect3 - 3*self.mean*self.stddev**2 - self.mean**3) / (self.stddev ** 3)
			self.hurst = (self.expect4 - 4*self.mean*self.expect3 + 6*self.mean**2*self.expect2 - 3*self.mean**4) / (self.stddev ** 4)

		# 更新中位数
		self.median = StatsPublic.calc_median(self.mean, self.stddev, self.skewness)
		# 更新众数
		self.mode = StatsPublic.calc_mode(self.mean, self.stddev, self.skewness)
		
class Distributions(Distribution):
	'''
	组合分布
	'''
	dist_class: type[Distribution] = None # type: ignore # 用于创建分布的类

	def __init__(self, lower_bound: float = None, upper_bound: float = None, dist_class: type[Distribution] = Distribution): # type: ignore
		super().__init__(lower_bound, upper_bound)
		self.dist_class: type[Distribution] = dist_class # 用于创建分布的类
		self.dists: list[Distribution] = [] # 用于存储多个分布的列表
		self.entropy: float = 0 # 熵值

	def get_info(self, label: str = '') -> str:
		'''
		打印信息
		params:
			label: str, 标签
		returns: str, 信息
		'''
		info = super().get_info(label) + '\n'
		info += f'【共有组{len(self.dists)}个:】\n'
		for i, dist in enumerate(self.dists):
			if dist == None:
				continue
			info += f"\t{dist.get_info(f'组{i}')},\n"
		info += '】'
		return info
	
	def _update_parameters(self):
		# 计算香农熵值
		self.entropy = 0
		for dist in self.dists:
			if dist == None or dist.n == 0:
				continue
			prob = dist.n / self.n
			self.entropy -= prob * np.log2(prob)

		return super()._update_parameters()
	
class GroupedDistributions(Distributions):
	'''
	固定分组分布
	'''

	def __init__(self, lower_bound:float = None , upper_bound: float = None, groups=10, dist_class: type[Distribution] = Distribution): # type: ignore
		super().__init__(lower_bound, upper_bound, dist_class)
		self.groups: int = groups # 分组数目
		self.group_size = (self.upper_bound - self.lower_bound) / self.groups # 每个组的范围
		self.dists = [self.dist_class(self.lower_bound + i * self.group_size, self.lower_bound + (i + 1) * self.group_size) for i in range(self.groups)] # 用于存储多个分布的列表


	def append_data(self, x: float, n: float = 1, N: float = None): # type: ignore
		N = self.n + n if N == None else N

		# 首先将数据添加到总体数据中
		super().append_data(x, n, N)

		# 分组统计
		group_size = (self.upper_bound - self.lower_bound) / self.groups # 每个组的范围

		group_index = (x - self.lower_bound) // group_size # 计算所在的组
		if group_index >= 0 and group_index < self.groups: # 范围内的组
			self.dists[int(group_index)].append_data(x, n, N) # 添加到相应组的分布中

	def calc_deviation(self, x: FloatOrArray, type: str = 'mean')->FloatOrArray:
		'''
		计算离散度
		'''
		for group in self.dists:
			if group.n == 0:
				continue
			return StatsPublic.calc_on_array_or_single(lambda x: group.calc_deviation(x, type) if group.in_range(x) else 0, x)
		return StatsPublic.calc_on_array_or_single(lambda x: super().calc_deviation(x, type), x)